from flask import Flask, request, jsonify
import jieba
import fasttext

app = Flask(__name__)

model = fasttext.load_model('./model/fasttext/black-news-classification.bin')
label2tag = {
    '__label__0': '保留',
    '__label__1': '舍弃'
}


@app.route('/check_black_news', methods=['POST'])
def check_black_news():
    text = request.form['text']
    text = ' '.join([w.strip() for w in jieba.cut(text) if len(w.strip()) > 0])
    labels, probs = model.predict(text, 2)
    res = {label2tag[label]: prob for label, prob in zip(labels, probs)}
    return jsonify(res)
